#!/usr/bin/env python3
import os
import cv2
import numpy as np
import time
import threading
import queue
from dataclasses import dataclass
from typing import Optional
import traceback
from opt_visualizer import MergedPosePositionVisualizer  # 如果使用优化的可视化器，请取消注释
# from opt_visualizer import RealtimeOptimizedVisualizer   # 如果使用实时优化的可视化器，请取消注释
from image_provider import ImageProvider
from seg_pose2 import SegPose
from compute_position import ComputePosition  
from utils import get_depth_image
from logger import setup_logger

logger = setup_logger()

@dataclass
class ImageData:
    """图像数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    color_path: Optional[str] = None

@dataclass
class ModelResult:
    """模型结果数据结构"""
    timestamp: str
    readable_timestamp: str
    color_img: np.ndarray
    depth_img: np.ndarray
    seg_result: any
    pose_result: any
    color_path: Optional[str] = None
    process_time: float = 0.0
    queue_wait_time: float = 0.0  # 新增：队列等待时间

@dataclass
class PositionResult:
    """位置计算结果数据结构"""
    timestamp: str
    readable_timestamp: str
    pose: any
    vector: any
    color_path: Optional[str] = None
    process_time: float = 0.0
    queue_wait_time: float = 0.0  # 新增：队列等待时间

def setup_standard_position(seg_pose, compute_position):
    """设置标准位置 - 独立函数"""
    st_img = cv2.imread("./data/Color/st.jpg")
    st_depth_img = get_depth_image("./data/Color/st.jpg")

    logger.info("==========设置标准位置==========")
    
    # 创建ImageData对象
    st_image_data = ImageData(
        timestamp=str(int(time.time() * 1000)),
        readable_timestamp="standard_image",
        color_img=st_img,
        depth_img=st_depth_img,
        color_path="./data/Color/st.jpg"
    )
    
    # 提交标准图像处理
    task_id = seg_pose.put(st_image_data)
    if task_id is None:
        logger.error("==========标准图片提交失败==========")
        return None, None
    
    # 等待结果
    max_wait_time = 30  # 最大等待30秒
    start_wait = time.time()
    
    while time.time() - start_wait < max_wait_time:
        seg_result, pose_result, returned_image_data = seg_pose.get_result()
        if seg_result is not None and pose_result is not None:
            logger.info("==========获取到标准图片的处理结果==========")
            break
        time.sleep(0.1)
    else:
        logger.error("==========等待标准图片处理结果超时==========")
        return None, None

    if seg_result is None or pose_result is None:
        logger.error("==========没有检测到标准图片的seg结果或pose结果==========")
        return None, None

    st_pose, st_vector = compute_position.set_standard(st_img, st_depth_img, seg_result, pose_result)
    if st_pose is None or st_vector is None:
        logger.error("==========标准位置设定失败==========")
        return None, None
    else:
        logger.info("==========标准位置设定成功==========")
        logger.info(f"标准位置: {st_pose}")
        logger.info(f"标准向量: {st_vector}")
        return st_pose, st_vector
    
class PerformanceMonitor:
    """性能监控类"""
    
    def __init__(self):
        self.lock = threading.Lock()
        self.stats = {
            'model_times': [],
            'position_times': [],
            'viz_times': [],
            'queue_wait_times': {
                'model': [],
                'position': [],
                'viz': []
            },
            'errors': {
                'model': 0,
                'position': 0,
                'viz': 0
            },
            'queue_drops': {  # 新增：记录队列丢弃的数据数量
                'image': 0,
                'model': 0,
                'position': 0
            }
        }
    
    def record_time(self, category: str, time_ms: float):
        """记录处理时间"""
        with self.lock:
            if category in self.stats:
                self.stats[category].append(time_ms)
    
    def record_queue_wait(self, stage: str, wait_time_ms: float):
        """记录队列等待时间"""
        with self.lock:
            if stage in self.stats['queue_wait_times']:
                self.stats['queue_wait_times'][stage].append(wait_time_ms)
    
    def record_error(self, category: str):
        """记录错误"""
        with self.lock:
            if category in self.stats['errors']:
                self.stats['errors'][category] += 1
    
    def record_queue_drop(self, queue_name: str, drop_count: int):
        """记录队列丢弃数量"""
        with self.lock:
            if queue_name in self.stats['queue_drops']:
                self.stats['queue_drops'][queue_name] += drop_count
    
    def get_summary(self):
        """获取统计摘要"""
        with self.lock:
            summary = {}
            for key, times in self.stats.items():
                if key == 'queue_wait_times':
                    summary[key] = {}
                    for stage, wait_times in times.items():
                        if wait_times:
                            summary[key][stage] = {
                                'avg': sum(wait_times) / len(wait_times),
                                'max': max(wait_times),
                                'min': min(wait_times),
                                'count': len(wait_times)
                            }
                elif key in ['errors', 'queue_drops']:
                    summary[key] = times
                elif times:  # 处理时间列表
                    summary[key] = {
                        'avg': sum(times) / len(times),
                        'max': max(times),
                        'min': min(times),
                        'count': len(times)
                    }
            return summary

class MultiThreadProcessor:
    """直接集成SegPose的优化处理器"""

    def __init__(self, st_pose, st_vector, seg_pose_instance, compute_position_instance, visualizer_instance,
                 max_queue_size: int = 10):
        self.st_pose = st_pose
        self.st_vector = st_vector
        
        # 只保留位置计算和可视化的队列，图像直接提交给seg_pose
        self.position_queue = queue.Queue(maxsize=max_queue_size)
        
        # 处理器
        self.seg_pose = seg_pose_instance
        self.compute_position = compute_position_instance
        
        # 使用实时优化的可视化器
        self.visualizer = visualizer_instance
        
        # 性能监控
        self.performance_monitor = PerformanceMonitor()
        
        # 线程控制
        self.stop_event = threading.Event()
        self.threads = []
        
        # 实时可视化统计
        self.viz_frame_count = 0
        self.viz_start_time = time.time()
        self.last_viz_stats_time = time.time()
        
        # 统计信息
        self.stats = {
            'images_submitted': 0,
            'models_processed': 0,
            'positions_processed': 0,
            'visualizations_processed': 0,
            'start_time': time.time(),
            'last_stats_time': time.time()
        }
        self.stats_lock = threading.Lock()
        
        logger.info(f"多线程主程序启动（直接SegPose模式）")

    def _model_worker(self):
        """模型结果获取工作线程 - 直接从seg_pose获取结果"""
        logger.info("模型结果获取线程已启动")
        
        consecutive_empty_count = 0
        max_consecutive_empty = 100  # 连续100次空结果后稍微增加等待时间
        
        while not self.stop_event.is_set():
            try:
                # 直接从seg_pose获取结果
                get_start = time.time()
                seg_result, pose_result, image_data = self.seg_pose.get_result()
                get_time = (time.time() - get_start) * 1000
                
                if seg_result is not None and pose_result is not None and image_data is not None:
                    # 获取到结果，重置空结果计数
                    consecutive_empty_count = 0
                    
                    logger.info(f"获取到模型结果: {image_data.readable_timestamp}")
                    
                    # 创建ModelResult对象
                    model_result = ModelResult(
                        timestamp=image_data.timestamp,
                        readable_timestamp=image_data.readable_timestamp,
                        color_img=image_data.color_img,
                        depth_img=image_data.depth_img,
                        seg_result=seg_result,
                        pose_result=pose_result,
                        color_path=image_data.color_path,
                        process_time=0.0,  # SegPose内部已处理
                        queue_wait_time=0.0
                    )
                    
                    # 尝试放入位置计算队列
                    success = self._safe_put(self.position_queue, model_result, "Position")
                    if success:
                        with self.stats_lock:
                            self.stats['models_processed'] += 1
                        
                        logger.debug(f"模型结果已送入位置计算队列: {image_data.readable_timestamp}")
                    
                else:
                    # 没有获取到结果
                    consecutive_empty_count += 1
                    
                    # 根据连续空结果次数调整等待时间
                    if consecutive_empty_count < max_consecutive_empty:
                        time.sleep(0.001)  # 1ms，保持高响应性
                    else:
                        time.sleep(0.01)   # 10ms，减少CPU占用
                
            except Exception as e:
                logger.error(f"模型结果获取错误: {e}")
                logger.error(f"错误详情: {traceback.format_exc()}")
                self.performance_monitor.record_error('model')
                time.sleep(0.1)  # 出错时等待一下
        
        logger.info("模型结果获取线程已退出")

    def _position_worker(self):
        """位置计算工作线程"""
        logger.info("位置计算线程已启动")
        
        while not self.stop_event.is_set():
            try:
                wait_start = time.time()
                model_result = self.position_queue.get(timeout=1.0)
                wait_time = (time.time() - wait_start) * 1000
                
                self.performance_monitor.record_queue_wait('position', wait_time)
                
                logger.debug(f"位置计算开始: {model_result.readable_timestamp}")
                
                start_time = time.time()
                pose, vector = self.compute_position.process_image_pair(
                    model_result.readable_timestamp,
                    model_result.color_img,
                    model_result.depth_img,
                    model_result.seg_result,
                    model_result.pose_result
                )
                process_time = (time.time() - start_time) * 1000
                
                self.performance_monitor.record_time('position_times', process_time)
                
                # 直接进行可视化，不再使用队列
                if pose is not None and vector is not None:
                    viz_start = time.time()
                    
                    # 使用实时优化的可视化
                    merged_image, frame_time = self.visualizer.visualize_merged(
                        pose,
                        self.st_pose,
                        self.st_vector,
                        vector,
                        model_result.color_path,
                        timestamp=model_result.readable_timestamp,
                        display=True,
                    )
                    
                    viz_time = (time.time() - viz_start) * 1000
                    self.performance_monitor.record_time('viz_times', viz_time)
                    
                    with self.stats_lock:
                        self.stats['positions_processed'] += 1
                        self.stats['visualizations_processed'] += 1
                    
                    self.viz_frame_count += 1
                    
                    logger.info(f"位置计算和可视化完成: {model_result.readable_timestamp}, "
                              f"位置计算: {process_time:.2f}ms, 可视化: {viz_time:.2f}ms")
                    logger.info(f"Pose: {pose}, Vector: {vector}")
                else:
                    logger.warning(f"位置计算失败: {model_result.readable_timestamp}")
                
            except queue.Empty:
                continue
            except Exception as e:
                logger.error(f"位置计算错误: {e}")
                logger.error(f"错误详情: {traceback.format_exc()}")
                self.performance_monitor.record_error('position')
        
        logger.info("位置计算线程已退出")

    def _monitor_worker(self):
        """增强的性能监控线程"""
        logger.info("增强性能监控线程已启动")
        
        while not self.stop_event.is_set():
            time.sleep(3.0)
            
            with self.stats_lock:
                current_time = time.time()
                elapsed_time = current_time - self.stats['start_time']
                
                if elapsed_time > 0:
                    submit_fps = self.stats['images_submitted'] / elapsed_time
                    process_fps = self.stats['models_processed'] / elapsed_time
                    perf_summary = self.performance_monitor.get_summary()
                    
                    logger.info("=== 增强性能诊断（直接SegPose模式）===")
                    logger.info(f"🔄 处理进度: 提交{self.stats['images_submitted']} → "
                               f"模型{self.stats['models_processed']} → "
                               f"位置{self.stats['positions_processed']} → "
                               f"可视化{self.stats['visualizations_processed']}")
                    logger.info(f"📊 提交FPS: {submit_fps:.2f}, 处理FPS: {process_fps:.2f}")
                    
                    # SegPose内部统计
                    segpose_stats = self.seg_pose.get_stats()
                    logger.info(f"📦 SegPose状态: 总put:{segpose_stats.get('total_put', 0)}, "
                               f"总complete:{segpose_stats.get('total_complete', 0)}, "
                               f"总get:{segpose_stats.get('total_get', 0)}, "
                               f"命中率:{segpose_stats.get('hit_rate', 0):.2f}")
                    
                    pending = self.seg_pose.get_pending_count()
                    logger.info(f"📦 SegPose队列: 处理中:{pending.get('processing', 0)}, "
                               f"已完成待取:{pending.get('completed_pending', 0)}")
                    
                    # 处理时间统计
                    for stage, key in [('位置计算', 'position_times'), ('可视化处理', 'viz_times')]:
                        if key in perf_summary:
                            stats = perf_summary[key]
                            logger.info(f"⏱️  {stage}: 平均{stats['avg']:.1f}ms, "
                                       f"最大{stats['max']:.1f}ms, 最小{stats['min']:.1f}ms")
                    
                    # 瓶颈分析和建议
                    position_backlog = self.position_queue.qsize()
                    
                    logger.info(f"📦 位置计算队列: {position_backlog}")
                    
                    # 智能瓶颈分析
                    if position_backlog > 5:
                        logger.warning("🚨 位置计算可能是瓶颈！")
                    elif pending.get('completed_pending', 0) > 10:
                        logger.warning("🚨 模型结果获取可能是瓶颈！")
                    elif position_backlog == 0 and pending.get('total_pending', 0) == 0:
                        logger.info("✅ 流水线运行完美，无瓶颈")
                    
                    logger.info("========================")
        
        logger.info("增强性能监控线程已退出")

    def start(self):
        """启动所有线程"""
        # 模型结果获取线程
        model_thread = threading.Thread(target=self._model_worker, name="ModelResultWorker")
        model_thread.start()
        self.threads.append(model_thread)
        
        # 位置计算线程（同时处理可视化）
        position_thread = threading.Thread(target=self._position_worker, name="PositionWorker")
        position_thread.start()
        self.threads.append(position_thread)
        
        # 增强的性能监控线程
        monitor_thread = threading.Thread(target=self._monitor_worker, name="MonitorWorker")
        monitor_thread.start()
        self.threads.append(monitor_thread)
        
        logger.info("所有优化线程已启动（直接SegPose模式）")

    def stop(self):
        """停止所有线程"""
        logger.info("正在停止所有线程...")
        self.stop_event.set()
        
        # 清理可视化器
        if hasattr(self, 'visualizer'):
            self.visualizer.cleanup()
        
        for thread in self.threads:
            thread.join(timeout=5.0)
            if thread.is_alive():
                logger.warning(f"线程 {thread.name} 未能在超时时间内停止")
        
        logger.info("所有线程已停止")

    def _safe_put(self, target_queue: queue.Queue, item, queue_name: str) -> bool:
        """安全地向队列添加元素"""
        try:
            target_queue.put_nowait(item)
            return True
        except queue.Full:
            # 队列满时，丢弃最老的元素
            try:
                target_queue.get_nowait()
                target_queue.put_nowait(item)
                self.performance_monitor.record_queue_drop(queue_name.lower(), 1)
                logger.warning(f"{queue_name}队列已满，丢弃了1个旧元素")
                return True
            except queue.Empty:
                logger.error(f"{queue_name}队列状态异常")
                return False

    def add_image(self, image_data: ImageData) -> bool:
        """直接将图像提交给SegPose处理"""
        task_id = self.seg_pose.put(image_data)
        success = task_id is not None
        
        if success:
            with self.stats_lock:
                self.stats['images_submitted'] += 1
            logger.debug(f"图像已提交给SegPose: {image_data.readable_timestamp}, task_id: {task_id}")
        else:
            logger.warning(f"图像提交失败: {image_data.readable_timestamp}")
        
        return success

# 修改主函数
def main():
    """优化的主函数（直接SegPose模式）"""
    seg_pose = SegPose()
    compute_position = ComputePosition()
    visualizer = MergedPosePositionVisualizer()
    
    # 设置标准位置
    st_pose, st_vector = setup_standard_position(seg_pose, compute_position)
    if st_pose is None or st_vector is None:
        logger.error("标准位置设定失败，程序退出")
        return

    # 配置参数
    from_folder = False
    max_queue_size = 5  # 实时模式下使用更小的队列
    
    # 创建图像提供器
    image_provider = ImageProvider(from_folder=from_folder)

    # 创建优化的处理器
    processor = MultiThreadProcessor(
        st_pose, st_vector, seg_pose, compute_position, visualizer,
        max_queue_size
    )
    
    try:
        processor.start()
        
        logger.info("开始处理图像（直接SegPose模式）...")
        
        while True:
            image_data = image_provider.get_next_image()
            if image_data is None:
                logger.info("所有图像读取完成")
                break
            
            success = processor.add_image(image_data)
            time.sleep(0.01)
            if not success:
                time.sleep(0.005)  # 实时模式下更短的等待时间
        
        # 等待队列处理完成
        logger.info("等待队列处理完成...")
        while True:
            pending = seg_pose.get_pending_count()
            position_backlog = processor.position_queue.qsize()
            
            if (pending.get('total_pending', 0) == 0 and position_backlog == 0):
                break
                
            time.sleep(0.5)
            logger.info(f"等待中... SegPose待处理: {pending.get('total_pending', 0)}, "
                       f"位置计算队列: {position_backlog}")
        
        logger.info("所有任务处理完成")
        
    except KeyboardInterrupt:
        logger.info("收到中断信号，正在关闭...")
    except Exception as e:
        logger.error(f"程序错误: {e}")
        logger.error(f"错误详情: {traceback.format_exc()}")
    finally:
        processor.stop()
        seg_pose.shutdown()  # 关闭SegPose
        image_provider.close()
        
        # 输出最终统计
        with processor.stats_lock:
            total_time = time.time() - processor.stats['start_time']
            logger.info("==========优化程序结束==========")
            logger.info(f"总共提交图像: {processor.stats['images_submitted']}")
            logger.info(f"总共处理图像: {processor.stats['models_processed']}")
            logger.info(f"总共可视化: {processor.stats['visualizations_processed']}")
            if processor.stats['images_submitted'] > 0:
                submit_fps = processor.stats['images_submitted'] / total_time
                process_fps = processor.stats['models_processed'] / total_time
                logger.info(f"平均提交帧率: {submit_fps:.2f} FPS")
                logger.info(f"平均处理帧率: {process_fps:.2f} FPS")
            
            # SegPose最终统计
            segpose_stats = seg_pose.get_stats()
            logger.info(f"SegPose最终统计: {segpose_stats}")
            
            # 可视化性能摘要
            viz_stats = processor.visualizer.get_performance_stats()
            if viz_stats:
                logger.info(f"可视化平均FPS: {viz_stats['estimated_fps']:.2f}")
                logger.info(f"可视化平均帧时间: {viz_stats['avg_frame_time_ms']:.1f}ms")
            
            logger.info(f"总运行时间: {total_time:.2f} 秒")


if __name__ == "__main__":
    main()